import torch
import torch.nn as nn
import config


class YoloV1Net(nn.Module):
    def __init__(self, S=config.S, B=config.B, C=config.C):
        super(YoloV1Net, self).__init__()
        self.B = B
        self.S = S
        self.C = C
        self.depth = self.B * 5 + self.C
        self.Conv = self.make_Conv()
        self.Classifier = self.make_Classifier()
        self.conv_network = nn.Sequential (*self.Conv)
        self.linear_network = nn.Sequential (*self.Classifier)

    def forward(self, x):
        """
        interfere the model
        :param x: the input data
        :return: the output of the network
        """
        out = self.conv_network (x)
        out = self.linear_network (out)
        out = torch.reshape(out, (x.size(0), self.S, self.S, self.depth))
        # the 0 dimension of x is the batch size of the input data
        return out

    def make_Conv(self):
        """
        generate the convolutional layers, the input will be (448, 448, 3)
        :return: the list of the conv_layers
        """
        layer1 = [
            nn.Conv2d(3, 64, 7, 2, 3),
            nn.LeakyReLU(negative_slope=0.1),
            nn.MaxPool2d(2, 2)
        ] # output : (112, 112, 64)
        layer2 = [
            nn.Conv2d(64, 192, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),
            nn.MaxPool2d(2, 2)
        ] # output : (56, 56, 192)
        layer3 = [
            nn.Conv2d(192, 128, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(128, 256, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(256, 256, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.MaxPool2d(2, 2)
        ] # output (28, 28, 512)
        layer4 = [
            nn.Conv2d(512, 256, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(512, 256, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(512, 256, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(512, 256, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(256, 512, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(512, 512, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(512, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.MaxPool2d(2, 2)
        ] # output : (14, 14, 1024)
        layer5 = [
            nn.Conv2d(1024, 512, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(512, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(1024, 512, 1, 1, 0),
            nn.LeakyReLU(negative_slope=0.1),
            nn.Conv2d(512, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            # The network layer in the paper
            nn.Conv2d(1024, 1024, 3, 2, 1),
            nn.LeakyReLU(negative_slope=0.1)

            # My layer of my realization
            # nn.MaxPool2d (2, 2)
        ] # output : (7, 7, 1024)
        layer6 = [
            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Conv2d(1024, 1024, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.1)
        ] # output : (7, 7, 1024)

        return layer1 + layer2 + layer3 + layer4 + layer5 + layer6

    def make_Classifier(self):
        """
        generate the layers of the full connect layer
        :return: the list of the full connect layer
        """
        out_features = self.S * self.S * self.depth
        layer = [
            nn.Flatten(),
            nn.Linear(self.S * self.S * 1024, 4096),
            nn.Dropout(),
            nn.LeakyReLU(negative_slope=0.1),

            nn.Linear(4096, out_features)
        ]
        
        return layer


if __name__ == '__main__' :
    net = YoloV1Net(config.S, config.B, config.C)
    print(net)